import pandas as pd
import baostock as bs
import numpy as np
import datetime
import MyTT
import click

from sqlalchemy import create_engine,text

engine = create_engine('mysql+pymysql://taomu:taomu-rcmu#2025@192.168.1.220:3306/taomu_stock')

def get_stock_history_data(dir,code,exists='append'):
    sdate = (datetime.datetime.today()-datetime.timedelta(days=600)).strftime("%Y-%m-%d")
    print(sdate,code)
    rs = bs.query_history_k_data_plus(code,"date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",start_date=sdate,end_date=None,frequency="d")
    data_list=[]
    while (rs.error_code == '0') & rs.next():
        row_data = rs.get_row_data()
        if row_data[-1] == "0":
            data_list.append(row_data)
    else:
        result = pd.DataFrame(data_list,columns=rs.fields)
        result['date'] = pd.to_datetime(result['date']);
        result['close'] = result['close'].astype(float);
        result['high'] = result['high'].astype(float);
        result['low'] = result['low'].astype(float);
        result['open'] = result['open'].astype(float);
        result['code'] = result['code'].astype(str)
        # 替换空字符串为 NaN
        result['volume'] = result['volume'].replace('', np.nan)
        # 转换为浮点数类型
        result['volume'] = pd.to_numeric(result['volume'], errors='coerce')
        # 填充缺失值为 0（可选）
        result['volume'] = result['volume'].fillna(0)
        result['volume'] = result['volume'].astype(float);
        result['preclose'] = result['preclose'].astype(float);
        # 替换空字符串为 NaN
        result['amount'] = result['amount'].replace('', np.nan)
        # 转换为浮点数类型
        result['amount'] = pd.to_numeric(result['amount'], errors='coerce')
        # 填充缺失值为 0（可选）
        result['amount'] = result['amount'].fillna(0)
        result['amount'] = result['amount'].astype(float); # 替换空字符串为 NaN
        result['turn'] = result['turn'].replace('', np.nan)
        # 转换为浮点数类型
        result['turn'] = pd.to_numeric(result['turn'], errors='coerce')
        # 填充缺失值为 0（可选）
        result['turn'] = result['turn'].fillna(0)
        result['turn'] = result['turn'].astype(float);
         # 替换空字符串为 NaN
        result['pctChg'] = result['pctChg'].replace('', np.nan)
        # 转换为浮点数类型
        result['pctChg'] = pd.to_numeric(result['pctChg'], errors='coerce')
        # 填充缺失值为 0（可选）
        result['pctChg'] = result['pctChg'].fillna(0)
        result['pctChg'] = result['pctChg'].astype(float);
        #  使用MyTT计算指标 5日均线、10日均线、20日均线、30日均线、60日均线 添加到result中 
        closes = result['close']
        result['obv'] = MyTT.OBV(closes,result['volume'])
        result['M5'] = MyTT.MA(closes,5)
        result['M10'] = MyTT.MA(closes,10)
        result['M20'] = MyTT.MA(closes,20)
        result['M30'] = MyTT.MA(closes,30)
        result['M50'] = MyTT.MA(closes,50)
        result['M60'] = MyTT.MA(closes,60)
        result['EMA5'] = MyTT.EMA(closes,5)
        result['EMA10'] = MyTT.EMA(closes,10)
        result['EMA20'] = MyTT.EMA(closes,20)
        result['BBI'] = MyTT.BBI(closes)
        result['boll_upper'], result['boll_middle'], result['boll_lower'] = MyTT.BOLL(closes)
        result['macd_dif'], result['macd_dea'], result['macd'] = MyTT.MACD(closes)
        result['rsi6'] = MyTT.RSI(closes,6)
        result['rsi12'] = MyTT.RSI(closes,12)
        result['k'],result['d'],result['j'] = MyTT.KDJ(result['high'],result['low'],closes)

        # 去掉开头60天数据
        result = result[result.index > 60]
        result.to_sql("taomu_stock_history",engine,if_exists=exists,index=False)
        #result.to_csv(dir+"/"+code+".csv", mode='w',header=True,index=False,method='multi')
    pass

def get_stock_code_list():
    d = 1
    df = bs.query_all_stock(day=datetime.date.today()-datetime.timedelta(days=d)).get_data()
    while(len(df) == 0):
        d += 1
        df = bs.query_all_stock(day=datetime.date.today()-datetime.timedelta(days=d)).get_data()
    else:
        print(len(df)) 
        if(len(df)>0):
            codes_df = df[((df['code'] < 'sz.300000') & (df['code'] > 'sh.689000') | (df['code'] < 'sh.688000') & (df['code'] > 'sh.009999')) & (~df['code_name'].str.startswith('ST')) & (~df['code_name'].str.startswith('*ST'))]
            return codes_df
    pass
    
@click.command()
@click.option("--dir", prompt="csv存放的目录", help="csv存放的目录")
def main(dir):
    # 清空表
    with engine.connect() as connection:
        try:
            connection.execute(text("delete from taomu_stock_list;"))
        except:
            pass
        try:
            connection.execute(text("delete from taomu_stock_history;"))
        except:
            pass
        connection.commit()
    bs.login()
    cpd = get_stock_code_list()
    cpd.to_sql("taomu_stock_list",engine,if_exists='append',index=False)
    clist = cpd['code'].to_list()
    get_stock_history_data(dir,clist.pop(),'replace')
    for code in clist:
        get_stock_history_data(dir,code)
    #get_stock_history_data("sh.600838")
    bs.logout()

if __name__ == '__main__':
    main()

